#include <Uefi.h>
#include <Library/UefiLib.h>
#include <Library/UefiApplicationEntryPoint.h>
#include <Library/BaseLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/PrintLib.h>
#include <Library/BaseCryptLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/UefiRuntimeServicesTableLib.h>
#include <Protocol/LoadedImage.h>
#include <Protocol/SimpleFileSystem.h>
#include <Guid/FileInfo.h>

// GUID for test variable namespace: {a1b2c3d4-1111-2222-3333-444455556666}
STATIC EFI_GUID mTestVarGuid = { 0xa1b2c3d4, 0x1111, 0x2222, {0x33,0x33,0x44,0x44,0x55,0x55,0x66,0x66} };

// Name used for the short-window test (1 minute)
#define FIRST_LOAD_TEST_VAR_NAME L"FirstLoadEpochTest"
#define ONE_MINUTE_SECONDS 60ULL

// Helper: convert EFI_TIME (assumed UTC) to epoch seconds (naive, supports typical ranges)
STATIC
UINT64
EfiTimeToEpochSeconds(
  IN CONST EFI_TIME *Time
  )
{
  // Simple algorithm: convert to days since 1970-01-01 then to seconds.
  // Not optimized; acceptable for a test helper.
  UINTN year = Time->Year;
  UINTN month = Time->Month;
  UINTN day = Time->Day;
  UINTN hour = Time->Hour;
  UINTN minute = Time->Minute;
  UINTN second = Time->Second;

  // Days per month for non-leap year
  CONST UINT8 DaysInMonth[] = { 31,28,31,30,31,30,31,31,30,31,30,31 };

  // Count days from 1970-01-01 to the given date (UTC). This is naive but fine for test.
  UINT64 days = 0;
  UINTN y;
  for (y = 1970; y < year; y++) {
    // leap year
    BOOLEAN leap = ((y % 4 == 0) && (y % 100 != 0)) || (y % 400 == 0);
    days += leap ? 366 : 365;
  }

  UINTN m;
  for (m = 1; m < month; m++) {
    days += DaysInMonth[m-1];
    // Feb leap adjust
    if (m == 2) {
      BOOLEAN leap = ((year % 4 == 0) && (year % 100 != 0)) || (year % 400 == 0);
      if (leap) {
        days += 1;
      }
    }
  }

  days += (day - 1);

  UINT64 secs = days * 86400ULL + (UINT64)hour * 3600ULL + (UINT64)minute * 60ULL + (UINT64)second;
  return secs;
}

// Test function: returns EFI_SUCCESS if within 1-minute window (or on first write), EFI_ACCESS_DENIED otherwise.
STATIC
EFI_STATUS
TestFirstLoadWindowOneMinute(
  IN EFI_HANDLE ImageHandle
  )
{
  EFI_STATUS Status;
  EFI_TIME Time;
  UINT64 CurrentEpoch;
  UINTN VarSize = sizeof(UINT64);
  UINT64 StoredEpoch;
  UINT32 Attributes;

  // Get current time
  Status = gRT->GetTime(&Time, NULL);
  if (EFI_ERROR(Status)) {
    Print(L"GetTime failed: %r\r\n", Status);
    return Status;
  }

  CurrentEpoch = EfiTimeToEpochSeconds(&Time);

  // Try read variable
  Status = gRT->GetVariable(FIRST_LOAD_TEST_VAR_NAME, &mTestVarGuid, &Attributes, &VarSize, &StoredEpoch);
  if (Status == EFI_NOT_FOUND) {
    // First run: write current epoch
    Attributes = EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS;
    Status = gRT->SetVariable(FIRST_LOAD_TEST_VAR_NAME, &mTestVarGuid, Attributes, sizeof(CurrentEpoch), &CurrentEpoch);
    if (EFI_ERROR(Status)) {
      Print(L"SetVariable failed: %r\r\n", Status);
      return Status;
    }
    Print(L"FirstLoadEpochTest written: %ld (epoch seconds). Allowing load.\r\n", CurrentEpoch);
    return EFI_SUCCESS;
  }

  if (EFI_ERROR(Status)) {
    Print(L"GetVariable failed: %r\r\n", Status);
    return Status;
  }

  // Compare
  if (CurrentEpoch < StoredEpoch) {
    Print(L"Current time is earlier than stored time (possible clock rollback). Denying.\r\n");
    return EFI_ACCESS_DENIED;
  }

  UINT64 Delta = CurrentEpoch - StoredEpoch;
  if (Delta <= ONE_MINUTE_SECONDS) {
    Print(L"Within 1-minute window (delta=%lds). Allowing load.\r\n", Delta);
    return EFI_SUCCESS;
  }

  Print(L"Outside 1-minute window (delta=%lds). Denying load.\r\n", Delta);
  return EFI_ACCESS_DENIED;
}

#define FIXED_INPUT_GUID_ASCII "732b8c84-0cb8-48a1-b525-8f032da2d885"

STATIC
CHAR16*
EncryptStringAndReturnHex (
  IN CONST CHAR8 *InputAscii
  )
{
  UINTN   InputLen;
  UINTN   PadLen;
  UINTN   TotalLen;
  UINT8  *Padded;
  UINT8  *Cipher;
  VOID   *AesCtx;
  UINTN   CtxSize;
  BOOLEAN Ok;
  UINTN   i;

  // 128-bit demo key and zero IV for illustration (DO NOT USE IN PRODUCTION)
  CONST UINT8 Key[16] = {
    0x00, 0x01, 0x02, 0x03,
    0x04, 0x05, 0x06, 0x07,
    0x08, 0x09, 0x0A, 0x0B,
    0x0C, 0x0D, 0x0E, 0x0F
  };
  UINT8 Iv[AES_BLOCK_SIZE];

  if (InputAscii == NULL) {
    return NULL;
  }

  SetMem(Iv, sizeof(Iv), 0);

  InputLen = AsciiStrLen(InputAscii);

  // PKCS#7 padding
  PadLen  = AES_BLOCK_SIZE - (InputLen % AES_BLOCK_SIZE);
  if (PadLen == 0) {
    PadLen = AES_BLOCK_SIZE;
  }
  TotalLen = InputLen + PadLen;

  Padded = AllocateZeroPool(TotalLen);
  if (Padded == NULL) {
    return NULL;
  }

  CopyMem(Padded, InputAscii, InputLen);
  for (i = InputLen; i < TotalLen; i++) {
    Padded[i] = (UINT8)PadLen;
  }

  Cipher = AllocateZeroPool(TotalLen);
  if (Cipher == NULL) {
    FreePool(Padded);
    return NULL;
  }

  CtxSize = AesGetContextSize();
  if (CtxSize == 0) {
    FreePool(Padded);
    FreePool(Cipher);
    return NULL;
  }

  AesCtx = AllocatePool(CtxSize);
  if (AesCtx == NULL) {
    FreePool(Padded);
    FreePool(Cipher);
    return NULL;
  }

  if (!AesInit(AesCtx, Key, 128)) {
    FreePool(AesCtx);
    FreePool(Padded);
    FreePool(Cipher);
    return NULL;
  }

  Ok = AesCbcEncrypt(AesCtx, Padded, TotalLen, Iv, Cipher);
  if (!Ok) {
    FreePool(AesCtx);
    FreePool(Padded);
    FreePool(Cipher);
    return NULL;
  }

  // Prepare hex string for ciphertext and return it (caller must FreePool)
  {
    // 2 hex chars per byte + null
    UINTN   HexChars = TotalLen * 2;
    CHAR16 *HexOut   = AllocateZeroPool((HexChars + 1) * sizeof(CHAR16));
    if (HexOut == NULL) {
      FreePool(AesCtx);
      FreePool(Padded);
      FreePool(Cipher);
      return NULL;
    }

    for (i = 0; i < TotalLen; i++) {
      // Write two hex chars at position i*2
      UnicodeSPrint(&HexOut[i * 2], (UINTN)((HexChars + 1) - (i * 2)) * sizeof(CHAR16), L"%02X", Cipher[i]);
    }

    FreePool(AesCtx);
    FreePool(Padded);
    FreePool(Cipher);
    return HexOut;
  }
}

/*
  Read all non-empty lines from an ASCII file on the same device the image was loaded from.
  Caller must FreePool each returned CHAR8* and FreePool the array.
  Returns NULL on error or if no lines read.
*/
STATIC
CHAR8 **
ReadLinesFromFile(
  IN EFI_HANDLE ImageHandle,
  IN CONST CHAR16 *FileName,
  OUT UINTN *OutCount
  )
{
  EFI_STATUS Status;
  EFI_LOADED_IMAGE_PROTOCOL *LoadedImage;
  EFI_SIMPLE_FILE_SYSTEM_PROTOCOL *SimpleFs;
  EFI_FILE_PROTOCOL *Root;
  EFI_FILE_PROTOCOL *File;
  EFI_FILE_INFO *FileInfo;
  UINTN InfoSize;
  UINTN FileSize;
  CHAR8 *FileBuf;
  UINTN ReadSize;
  UINTN i;
  UINTN LineCount = 0;
  CHAR8 **Lines = NULL;

  if (OutCount == NULL) {
    return NULL;
  }

  *OutCount = 0;

  Status = gBS->HandleProtocol(ImageHandle, &gEfiLoadedImageProtocolGuid, (VOID**)&LoadedImage);
  if (EFI_ERROR(Status)) {
    return NULL;
  }

  Status = gBS->HandleProtocol(LoadedImage->DeviceHandle, &gEfiSimpleFileSystemProtocolGuid, (VOID**)&SimpleFs);
  if (EFI_ERROR(Status)) {
    return NULL;
  }

  Status = SimpleFs->OpenVolume(SimpleFs, &Root);
  if (EFI_ERROR(Status)) {
    return NULL;
  }

  Status = Root->Open(Root, &File, (CHAR16*)FileName, EFI_FILE_MODE_READ, 0);
  if (EFI_ERROR(Status)) {
    return NULL;
  }

  InfoSize = 0;
  FileInfo = NULL;
  Status = File->GetInfo(File, &gEfiFileInfoGuid, &InfoSize, FileInfo);
  if (Status == EFI_BUFFER_TOO_SMALL) {
    FileInfo = AllocatePool(InfoSize);
    if (FileInfo == NULL) {
      File->Close(File);
      return NULL;
    }
    Status = File->GetInfo(File, &gEfiFileInfoGuid, &InfoSize, FileInfo);
  }
  if (EFI_ERROR(Status)) {
    if (FileInfo != NULL) {
      FreePool(FileInfo);
    }
    File->Close(File);
    return NULL;
  }

  FileSize = (UINTN)FileInfo->FileSize;
  FreePool(FileInfo);

  FileBuf = AllocatePool(FileSize + 1);
  if (FileBuf == NULL) {
    File->Close(File);
    return NULL;
  }

  ReadSize = FileSize;
  Status = File->Read(File, &ReadSize, FileBuf);
  File->Close(File);
  if (EFI_ERROR(Status) || ReadSize == 0) {
    FreePool(FileBuf);
    return NULL;
  }

  // Normalize buffer: treat CR and LF as separators
  // First pass: count non-empty lines
  {
    BOOLEAN InLine = FALSE;
    for (i = 0; i < ReadSize; i++) {
      CHAR8 c = FileBuf[i];
      if (c != '\r' && c != '\n') {
        if (!InLine) {
          LineCount++;
          InLine = TRUE;
        }
      } else {
        InLine = FALSE;
      }
    }
  }

  if (LineCount == 0) {
    FreePool(FileBuf);
    return NULL;
  }

  Lines = AllocateZeroPool(LineCount * sizeof(CHAR8*));
  if (Lines == NULL) {
    FreePool(FileBuf);
    return NULL;
  }

  // Second pass: extract lines
  {
    UINTN idx = 0;
    UINTN pos = 0;
    while (pos < ReadSize) {
      // skip separators
      while (pos < ReadSize && (FileBuf[pos] == '\r' || FileBuf[pos] == '\n')) {
        pos++;
      }
      if (pos >= ReadSize) {
        break;
      }
      UINTN start = pos;
      while (pos < ReadSize && FileBuf[pos] != '\r' && FileBuf[pos] != '\n') {
        pos++;
      }
      UINTN len = pos - start;
      CHAR8 *Line = AllocateZeroPool(len + 1);
      if (Line == NULL) {
        // cleanup
        for (i = 0; i < idx; i++) {
          FreePool(Lines[i]);
        }
        FreePool(Lines);
        FreePool(FileBuf);
        return NULL;
      }
      CopyMem(Line, &FileBuf[start], len);
      Line[len] = '\0';
      Lines[idx++] = Line;
    }
  }

  FreePool(FileBuf);
  *OutCount = LineCount;
  return Lines;
}

/*
  Write provided ASCII lines to a file (create or replace) on same device image loaded from.
  Caller-owned lines are not freed by this function.
*/
STATIC
EFI_STATUS
WriteLinesToFile(
  IN EFI_HANDLE ImageHandle,
  IN CONST CHAR16 *FileName,
  IN CHAR8 **Lines,
  IN UINTN Count
  )
{
  EFI_STATUS Status;
  EFI_LOADED_IMAGE_PROTOCOL *LoadedImage;
  EFI_SIMPLE_FILE_SYSTEM_PROTOCOL *SimpleFs;
  EFI_FILE_PROTOCOL *Root;
  EFI_FILE_PROTOCOL *File;
  UINTN i;

  if (Lines == NULL || Count == 0) {
    return EFI_INVALID_PARAMETER;
  }

  Status = gBS->HandleProtocol(ImageHandle, &gEfiLoadedImageProtocolGuid, (VOID**)&LoadedImage);
  if (EFI_ERROR(Status)) {
    return Status;
  }

  Status = gBS->HandleProtocol(LoadedImage->DeviceHandle, &gEfiSimpleFileSystemProtocolGuid, (VOID**)&SimpleFs);
  if (EFI_ERROR(Status)) {
    return Status;
  }

  Status = SimpleFs->OpenVolume(SimpleFs, &Root);
  if (EFI_ERROR(Status)) {
    return Status;
  }

  // If file exists, delete it first so we start fresh
  Status = Root->Open(Root, &File, (CHAR16*)FileName, EFI_FILE_MODE_READ, 0);
  if (!EFI_ERROR(Status)) {
    // Delete closes file
    File->Delete(File);
  }

  // Create new file
  Status = Root->Open(Root, &File, (CHAR16*)FileName, EFI_FILE_MODE_READ | EFI_FILE_MODE_WRITE | EFI_FILE_MODE_CREATE, 0);
  if (EFI_ERROR(Status)) {
    return Status;
  }

  for (i = 0; i < Count; i++) {
    if (Lines[i] == NULL) {
      continue;
    }
    UINTN WriteSize = AsciiStrLen(Lines[i]);
    // write the line bytes
    Status = File->Write(File, &WriteSize, (VOID*)Lines[i]);
    if (EFI_ERROR(Status)) {
      File->Close(File);
      return Status;
    }
  }

  File->Close(File);
  return EFI_SUCCESS;
}

EFI_STATUS
EFIAPI
UefiMain (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE *SystemTable
  )
{
  EFI_STATUS Status;
  // Run 1-minute first-load test for quick validation
  Status = TestFirstLoadWindowOneMinute(ImageHandle);
  if (EFI_ERROR(Status)) {
    Print(L"First-load 1-minute test failed or denied: %r\r\n", Status);
  } else {
		Print(L"First-load 1-minute test passed, proceeding to file processing...\r\n");
	}

	return EFI_SUCCESS;

  UINTN Count = 0;
  CHAR8 **Lines = NULL;
  UINTN i;

  // Try read InputString.list from the device the image was loaded from
  Lines = ReadLinesFromFile(ImageHandle, L"InputString.list", &Count);
  if (Lines == NULL || Count == 0) {
    // Fallback: process the fixed input
		Print(L"InputString.list not found or empty, processing fixed input\r\n");
    CHAR16 *Hex = EncryptStringAndReturnHex(FIXED_INPUT_GUID_ASCII);
    if (Hex == NULL) {
      Print(L"Encryption failed or returned no data\r\n");
      return EFI_SUCCESS;
    }
    Print(L"Input (ASCII): %a\r\n", FIXED_INPUT_GUID_ASCII);
    Print(L"Encrypted (AES-CBC/PKCS7) hex: %s\r\n", Hex);
    FreePool(Hex);
    return EFI_SUCCESS;
  }

  // Build output ASCII lines in format: original:encrypted\n
  CHAR8 **OutLines = AllocateZeroPool(Count * sizeof(CHAR8*));
  if (OutLines == NULL) {
    // cleanup
		Print(L"Failed to allocate OutLines\r\n");
    for (i = 0; i < Count; i++) {
      FreePool(Lines[i]);
    }
    FreePool(Lines);
    Print(L"Allocate OutLines failed\r\n");
    return EFI_SUCCESS;
  }

  for (i = 0; i < Count; i++) {
    CHAR8 *Orig = Lines[i];
    CHAR16 *Hex = EncryptStringAndReturnHex(Orig);
    if (Hex == NULL) {
      // write original with empty hex
			Print(L"Encryption failed for line %u, writing original only\r\n", i);
      UINTN Len = AsciiStrLen(Orig) + 2; // for ':' and '\n' and terminating +1 handled in Allocate
      CHAR8 *Entry = AllocateZeroPool(Len + 1);
      if (Entry != NULL) {
        CopyMem(Entry, Orig, AsciiStrLen(Orig));
        Entry[AsciiStrLen(Orig)] = ':';
        Entry[AsciiStrLen(Orig) + 1] = '\n';
        Entry[AsciiStrLen(Orig) + 2] = '\0';
      }
      OutLines[i] = Entry;
    } else {
      // Convert Hex (CHAR16) to ASCII
      UINTN HexLen = StrLen(Hex);
      UINTN OrigLen = AsciiStrLen(Orig);
      // overall len: OrigLen + 1(':') + HexLen + 1('\n')
      UINTN Total = OrigLen + 1 + HexLen + 1;
      CHAR8 *Entry = AllocateZeroPool(Total + 1);
      if (Entry != NULL) {
        // copy original
        CopyMem(Entry, Orig, OrigLen);
        Entry[OrigLen] = ':';
        // convert Hex (CHAR16) to ASCII, assuming only ASCII hex chars
        UINTN j;
        for (j = 0; j < HexLen; j++) {
          CHAR16 c = Hex[j];
          Entry[OrigLen + 1 + j] = (CHAR8)(c & 0xFF);
        }
        Entry[OrigLen + 1 + HexLen] = '\n';
        Entry[OrigLen + 1 + HexLen + 1] = '\0';
      }
      OutLines[i] = Entry;
      FreePool(Hex);
    }
  }

  // Write to OutputString.list
  Status = WriteLinesToFile(ImageHandle, L"OutputString.list", OutLines, Count);
  if (EFI_ERROR(Status)) {
    Print(L"Failed to write OutputString.list: %r\r\n", Status);
  } else {
    Print(L"Wrote %u lines to OutputString.list\r\n", Count);
  }

  // cleanup
  for (i = 0; i < Count; i++) {
    if (Lines[i] != NULL) {
      FreePool(Lines[i]);
    }
    if (OutLines[i] != NULL) {
      FreePool(OutLines[i]);
    }
  }
  FreePool(Lines);
  FreePool(OutLines);

  return EFI_SUCCESS;
}
